#!/usr/bin/python
#coding:utf-8

# Copyright 2011 Nicolau Leal Werneck, Anna Helena Reali Costa and
# Universidade de São Paulo
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

###############################################################################
### This program uses the camori module to find the camera orientation relative
### to a "natural" reference frame in the photographed envronment, which is
### defined by the restricted and mutually-orthogonal directions to where all
### the environment edges are pointing to. The input shoul dbe a normal camera,
### following the pinhole model.

import sys

import time

import matplotlib
#matplotlib.use('Agg')

from pylab import *

import scipy.io

from numpy import dot

from camori import PicturePinhole, quaternion_to_matrix, measure_error, quaternion_product

import Image


if __name__ == '__main__':

  ## Plot stuff immediately
  ion()

  rc('text',usetex=False)

  #dir_colors=['#aa00aa', '#00aa00', '#0000aa']
  dir_colors=['#a370ff', '#51c373', '#ea6949'] #mygucci

  ## Avoid zero divide warnins...
  np.seterr(divide='ignore')

  #Downscale = True
  Downscale = False
  PlotStuff=True
  #PlotStuff=False

  UsePIL = True

  #################################################################
  ## Load image and initialize pic object

  ## Sets filename from input argument
  if len(sys.argv)<2:
    print sys.argv[0], '<filename> <grid_spacing> <gradient_threshold> <angle_variance> <smoothing_factor>'
    exit

  filename=sys.argv[1]

  ## Opens the image and then optionally smooth out and downscale.
  im = Image.open (filename)
  frame = array(im, dtype=float)
  imr = array(im, dtype=float)
  imr = imr[:,:,:3] #remove alpha channel
  # Smooth out
  if len(sys.argv) > 5 and double(sys.argv[5])>0:
    for c in range(3):
      imr[:,:,c] = scipy.ndimage.gaussian_filter( imr[:,:,c], double(sys.argv[5]))

  ## Rescales image
  if Downscale:
    zzz = zeros( (imr.shape[0]/4,imr.shape[1]/4,3) )
    for chan in range(3):
      zzz[:,:,chan]=scipy.ndimage.interpolation.zoom(imr[:,:,chan], 0.25)
  else:
    zzz = zeros( (imr.shape[0]/1,imr.shape[1]/1,3) )
    for chan in range(3):
      zzz[:,:,chan]=scipy.ndimage.interpolation.zoom(imr[:,:,chan], 1.0)


  ## === Intrinsic parameters ===
  ## This value should definitely come from some serious camera calibration
  ## procedure. it's being set here just based on a hunch.
  #focal_distance = 600

  ## For a 28mm equivalent:
  focal_distance = zzz.shape[1] * 28./36.
  ## The principal point (projection center) is assumed to be the
  ## middle of the image... Could be set.
  ################################################################

  ## Creates picture object
  pic = PicturePinhole(zzz, focal_distance)

  ## Don take last columns in consideration...
  pic.mask[:,-3:] = 0 #remove last column


  ## Run the grid sweep, and generate the pic.edgels list
  ## Edgel extractor parameters
  gspc=int(sys.argv[2])#10
  glim=float(sys.argv[3]) ** -1 #35
  p3=float(sys.argv[4])#0.25
  pic.p3=p3
  time1=time.clock()
  pic.extract_edgels(gspc, glim)

  time1=time.clock()-time1
  print 'Time sweeping:',time1
  ##
  #################################################################



  #################################################################
  ## Execute the optimizaiton
  if True:
  #if False:
    # the_inis = array([ [0, 0, .10], [0, 0, -.10],
    #    [.10, 0, 0], [-.10, 0, 0], [0, .10, 0], [0, -.10, 0],])
    # the_inis = array([ [0, 0, .20], [0, 0, -.20],
    #    [.20, 0, 0], [-.20, 0, 0], [0, .20, 0], [0, -.20, 0],])
    the_inis = array([ [0.04,-.2,0], [0.04,-.25, 0], [.04,-.15, 0],])

    time1=time.clock()
    optims = [ pic.find_orientation(xini) for xini in the_inis]
    time1=time.clock()-time1
    print 'Time optimizing:',time1
    mydic = dict( [(double(x[1]),x[0]) for x in optims ] )
    xopt=mydic[min(mydic)]
  ##
  #################################################################

  else:
    #222155
    #xopt = array([ 0.03505814,  0.58781221,  0.09545219])
    #xopt = array([ 0.08802025, -0.117076,    0.04462472]) 
    xopt = array([0.04,-.143,0])

  pic.orientation=xopt
  print 'Ans:', xopt


  #################################################################
  ## Plot the edgels and vanishing point directions

  ## Make sure we have the right labels
  pic.calculate_MAP(quaternion_to_matrix(pic.orientation))

  if PlotStuff:
    figure(1, figsize=(12.8,7.2))
    title('[SteetSLAM] Edgels detected over grid mask, labeled')
    imshow(pic.frame/260., interpolation='nearest')
    kk=0.40*gspc
    labs=['x','y','z']
    for lab,col in enumerate(dir_colors):
      ed = pic.edgels[pic.labels==lab]

      th_c = ed[:,0]
      ph_c = ed[:,1]

      plot((c_[th_c,th_c] + kk*c_[-ed[:,3],+ed[:,3]]).T,
         (c_[ph_c,ph_c] + kk*c_[+ed[:,2],-ed[:,2]]).T,
         '-', color=col, lw=2, label=labs[lab])
    axis('equal')

    savefig('the_edgels.png', dpi=100)

    figure(2, figsize=(12.8,7.2))
    title('[SteetSLAM] Predicted edge directions at multiple points')
    imshow(pic.frame/260., interpolation='nearest', extent = (-pic.Iwidth/2,pic.Iwidth/2,pic.Iheight/2,-pic.Iheight/2,))
    pic.plot_vdirs(gca(), 36, quaternion_to_matrix(pic.orientation))
    axis('equal')

    savefig('the_vps.png', dpi=100)
  ##
  #################################################################

  R = quaternion_to_matrix(pic.orientation)

  outsize = 500.
  os = 1.0
  outfd = 40000.
  ## Find the limits of rectified image
  # pp = array([
  #     [-0.5*outsize,-0.5*outsize, outfd],
  #     [-0.5*outsize, os*outsize, outfd],
  #     [ 2*outsize, os*outsize, outfd],
  #     [ 2*outsize,-0.5*outsize, outfd],
  #     ])
  pp = array([
      [outfd, -.75*outfd, 6.5*outfd],
      [outfd,  .75*outfd, 6.5*outfd],
      [outfd,  .75*outfd, 1.5*outfd],
      [outfd, -.75*outfd, 1.5*outfd],
      ])

  qq = dot(pp,R)

  q = qq[:,:2]*focal_distance/qq[:,[2,2]] * (4 if Downscale else 1)
  q += array([2000,1500])

  if UsePIL:
    aa = im.transform( ( 5.*outsize, 1.5*outsize),
                       Image.QUAD, q.ravel(), Image.BILINEAR )
    # aa = im.transform( ( outsize, outsize),
    #                    Image.QUAD, [0,0,0,3000,4000,3000,4000,0], Image.BILINEAR )

    pixels = aa.load()
    #pixels[round(omiddlex),round(omiddley)]=(255,0,0)
    pixels[10,10]=(255,0,0)

    aa.save(filename + '-RECTIFIED.png')
